Custom WebSocket client

Overview

The SilentShardWebsocketClient class is a protocol-aware WebSocket client designed specifically for MPC (Multi-Party Computation) operations. It serves as a thin wrapper around the WebsocketClient which implements NetworkClient from com.silencelaboratories.network module.

1. SilentShardWebsocketClient implements WebsocketClient

• Responsibility:
  • Raw WebSocket connection management
  • Message queueing and asynchronous I/O
  • Binary/text data transmission
• Customization Use Cases:
  • Add Action specific logic for example perform some logic on keygen connect
  • Replace WebSocket library (e.g., Socket.IO, different WebSocket impl)
  • Modify low-level message queueing logic

Example implementing custom websocket client

App.tsx

package com.silencelaboratories.network.silentshard

import com.silencelaboratories.network.httpClient
import com.silencelaboratories.network.websocket.WebsocketClient
import com.silencelaboratories.network.websocket.WebsocketConfig
import io.ktor.client.plugins.websocket.DefaultClientWebSocketSession
import io.ktor.client.plugins.websocket.WebSockets
import io.ktor.client.plugins.websocket.webSocketSession
import io.ktor.client.request.HttpRequestBuilder
import io.ktor.client.request.header
import io.ktor.client.request.headers
import io.ktor.client.request.post
import io.ktor.client.request.setBody
import io.ktor.client.statement.HttpResponse
import io.ktor.client.statement.readRawBytes
import io.ktor.http.ContentType
import io.ktor.http.HttpHeaders
import io.ktor.http.contentType
import io.ktor.http.takeFrom
import io.ktor.websocket.Frame
import io.ktor.websocket.readReason
import kotlinx.coroutines.withTimeout


/**
 * Default [WebsocketClient] for SilentShard.
 *
 * This class could be extended to add custom logic before or after any websocketClient
 * operations such as connecting, read, etc. Or users are open to use their own version of
 * websocket client by extending [WebsocketClient].
 * */

open class SilentShardWebsocketClient<T : SilentShardWebsocketAction>(
    private val websocketConfig: WebsocketConfig,
) : WebsocketClient<T> {
    private val tag = "SilentShardWebsocket"

    private val httpClient = httpClient {
        install(WebSockets)
        /*install(Logging) {
            logger = object : Logger {
                override fun log(message: String) {
                    Log.d("Logger Ktor =>", message)
                }
            }
            level = LogLevel.ALL
        }*/
    }

    private var session: DefaultClientWebSocketSession? = null

    /**
     * Perform websocket connection with another party.
     * @param action The mpc action to perform. For example, Keygen, Refresh, etc.
     * @param params Additional parameters for the mpc action. For example, KeyId in KeyRefresh action, etc.
     * */
    override suspend fun connect(action: T, params: String?) {
        when (action) {
            is SilentShardWebsocketAction.DuoAction.Keygen -> {
                /*
                * Perform additional stuff for keygen. */
            }
            is SilentShardWebsocketAction.DuoAction.Import -> TODO()
            is SilentShardWebsocketAction.DuoAction.KeyRefresh -> TODO()
            is SilentShardWebsocketAction.DuoAction.Signature -> TODO()
            is SilentShardWebsocketAction.DuoAction.Reconcile -> TODO()
            is SilentShardWebsocketAction.TrioAction.Export -> TODO()
            is SilentShardWebsocketAction.TrioAction.Import -> TODO()
            is SilentShardWebsocketAction.TrioAction.KeyRefresh -> TODO()
            is SilentShardWebsocketAction.TrioAction.Keygen -> TODO()
            is SilentShardWebsocketAction.TrioAction.PreSignature -> TODO()
            is SilentShardWebsocketAction.TrioAction.PreSignatureFinal -> TODO()
            is SilentShardWebsocketAction.TrioAction.Reconcile -> TODO()
            is SilentShardWebsocketAction.TrioAction.Recovery -> TODO()
            is SilentShardWebsocketAction.TrioAction.Signature -> TODO()
        }
        val url = buildString {
            append(websocketConfig.websocketUrl)
            append(action.buildWebsocketConnectionUrl())
            params?.let { append("/$it") }
        }
        log("Connecting to $url")
        var request: HttpRequestBuilder? = null
        val requestBuilder: HttpRequestBuilder.() -> Unit = {
            url { takeFrom(url) }
            header(HttpHeaders.SecWebSocketProtocol, action.protocol)
            request = this
        }
        session = httpClient.webSocketSession(requestBuilder)
        log(request?.body.toString() ?: "")
        log("Connected")
        websocketConfig.authenticationToken?.let {
            log("Sending auth token")
            session?.send(Frame.Text(it))
            log("Auth token sent")
        }
    }

    /**
     * Send mpc bytes to another party.
     *
     * @param bytes The bytes to send.
     */
    override suspend fun send(bytes: ByteArray) {
        session?.let {
            log("Sending bytes: ${bytes.size}")
            it.send(Frame.Binary(true, bytes))
            log("Bytes sent")
        } ?: log("Session is null, cannot send bytes")
    }

    /**
     * Send mpc text to another party.
     *
     * @param text The text to send.
     * */
    override suspend fun send(text: String) {
        session?.let {
            log("Sending text: $text")
            it.send(Frame.Text(text))
            log("Text sent")
        } ?: log("Session is null, cannot send text")
    }

    /**
     * Read mpc bytes from another party.*/
    override suspend fun read(): ByteArray {
        return session?.let { wsSession ->
            log("Reading from WebSocket")
            when (val frame = wsSession.incoming.receive()) {
                is Frame.Binary -> frame.data.also { log("Received bytes: ${it.size}") }
                is Frame.Text -> frame.data.also { log("Received text: ${it.decodeToString()}") }
                is Frame.Close -> {
                    log("Received close frame: ${frame.readReason()}")
                    disconnect()
                    ByteArray(0)
                }

                else -> ByteArray(0)
            }
        } ?: run {
            log("Session is null, cannot read")
            ByteArray(0)
        }
    }

    /**
     * Perform websocket disconnection with another party.
     * */
    override suspend fun disconnect() = run {

        session?.let { wsSession ->
            log("Initiating disconnection from WebSocket")
            log("Sending close frame")

            try {
                wsSession.send(Frame.Close())
            } catch (exception: Exception) {
                log("Error sending close frame : ${exception.message}")
            }

            val closeReason = withTimeout(5000) {
                wsSession.closeReason.await()
            }

            when (closeReason?.code) {

                null -> {
                    throw Exception("WebSocket Close Reason : Unknown")
                }

                else -> {
                    log("WebSocket Close Reason : ${closeReason.knownReason} : ${closeReason.code} : message : ${closeReason.message}")
                }
            }
            session = null
        } ?: log("Session is null, no disconnection needed")

    }


    /**
     * Send a post-request to another party.
     *
     * @param action The mpc action to perform requires a post-request.
     * @param authHeader The authentication header to send with the post-request.
     * @param body The request body tobe sent.
     * It could be a JSON string or ByteArray.
     * Child class
     * should perform a type check and set `contentType` accordingly.
     * For example, if body is a JSON string - `contentType(ContentType.Application.Json)`.
     * If body is ByteArray - `contentType(ContentType.Application.OctetStream)`.
     * */
    override suspend fun sendPostRequest(
        action: T,
        authHeader: String?,
        body: Any,
    ): ByteArray {
        val url = buildString {
            append(websocketConfig.postRequestUrl)
            append(action.buildPostRequestUrl())
        }
        log("Sending POST request to $url")
        val response: HttpResponse = httpClient.post(url) {
            headers {
                authHeader?.let {
                    log("Adding auth header")
                    header(HttpHeaders.Authorization, "Bearer $it")
                }
            }
            when (body) {
                is String -> contentType(ContentType.Application.Json)
                is ByteArray -> contentType(ContentType.Application.OctetStream)
            }
            setBody(body)
        }
        return response.readRawBytes()
    }

    private fun log(message: String) {
        com.silencelaboratories.network.log(tag, message)
    }
}